import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain_core.prompts import StringPromptTemplate
from langchain_openai import ChatOpenAI
import inspect

# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_BASE_URL")
model = os.getenv("DEFAULT_MODEL")

llm = ChatOpenAI(
    model=model,
    temperature=0,
    #timeout=None,
    max_retries=2,
    api_key=api_key,
    base_url=base_url
)

# 定义一个简单的函数作为示例效果
def hello_world(abc):
    print("Hello, world!")
    return abc

PROMPT = """
你是一个非常有经验和天赋的程序员，现在给你如下函数名称，你会按照如下格式，输出这段代码的名称、源代码、中文解释。
函数名称: {function_name}
源代码: {source_code}
代码解释:
"""

def get_source_code(function_name):
    #获得源代码
    return inspect.getsource(function_name)


# 自定义的模板class
class CustomPrompt(StringPromptTemplate):

    def format(self, **kwargs) -> str:
        # 获得源代码
        source_code = get_source_code(kwargs["function_name"])

        # 生成提示词模板
        prompt = PROMPT.format(
            function_name=kwargs["function_name"].__name__,
            source_code=source_code
        )
        return prompt


a = CustomPrompt(input_variables=["function_name"])
pm = a.format(function_name=hello_world)

res = llm.invoke(pm)
print(res.content)